{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Improving network performance on unbalanced data via augmentation 🦇\n",
    "By Jonas Bengt Carina Håkansson\n",
    "\n",
    " Part of the 2022 AI Residency Program"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "Read the [DeepLabCut Blog post](https://deeplabcut.medium.com/balancing-bats-d3c5bcfb71e1) about how one can solve unbalanced data with flipping-keypoint augmentation, then see below for the process and code to do so!\n",
    "\n",
    "## The dataset\n",
    "This dataset consists of videos of bats flying freely at the [Austin Bat Refuge](https://austinbatrefuge.org/). The videos were shot by Dr. Aaron Corcoran, Head of the [Sensory & Aerial Ecology lab](https://sonarjamming.com/) at the University of Colorado, Colorado Sprungs (UCCS) in 2020. The dataset is comprised of 26 video triplets. Each triplet consists of three videos depicting the same bat flight from different angles.\n",
    "\n",
    "The videos have been partially digitized in [DLTdv](https://github.com/tlhedrick/dltdv) after which custom MATLAB code has been used to extract the digitized frames and coordinates in a format that DeepLabCut can read.\n",
    "\n",
    "## Background\n",
    "This text documents my experience in using data augmentation to improve the accuracy of DeepLabCut in tracking bodyparts of the bats in the videos during my time as a [DeepLabCut AI Resident](https://www.deeplabcutairesidency.org/) during the summer of 2022."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the DeepLabCut project\n",
    "We start by defining the path we will be working in. Then we create the project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false,
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "\n",
    "project_folder = \"/home/user/projects/\" #the folder in which the DLC project will be created\n",
    "deeplabcut.create_new_project(project='bat_augmentation_austin_2020_bat_data',experimenter='DLC',videos=['/home/user/dummyVideos/'],working_directory=project_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Edit the config file\n",
    "\n",
    "Now that the project is created we need to modify our project config file. Open config.yaml in the project folder.\n",
    "\n",
    "The config file can be modified with any text editor, but I will do it via Python commands to highlight the possibility of automating this process.\n",
    "\n",
    "![what to change in the config file](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950854851-BEJVFJYTWOZDXFFF1091/config001_videos_and_bodyparts.png?format=750w)\n",
    "\n",
    "*The list of videos and the list of bodyparts need to change to reflect the actual videos and bodyparts*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "#define config file\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "\n",
    "#import tools for modifying our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config, edit_config\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Bat bodyparts](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950854711-GE81YU6I16GB4EXYLGWG/bat_markers.png?format=750w)\n",
    "*The locations of the labels/bodyparts on the bats.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# replace the default list of bodyparts with a list of the parts that we have actually digitized\n",
    "edit_config(config_path,{\"bodyparts\":['t3L', 'wstL', 't5L', 'elbL', 'shdL', 'ankL', 'nl', 'str', 'lmb', 'shdR', 'ankR', 'elbR', 'wstR', 't5R', 't3R', 'tail']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "#fetch the list of videos from an older project using the same videos\n",
    "videolist = read_config(\"/home/user/projects/old_project-DLC-2022-08-03/config.yaml\")[\"video_sets\"]\n",
    "edit_config(config_path,{'video_sets':videolist})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert training data to DLC format\n",
    "Since this data was labeled outside of DeepLabCut, we first used our own code to convert it into csv formatted csv files, we then use the function ```convertcsv2h5``` to convert these csv:s into h5 files that DeepLabCut can read.\n",
    "\n",
    "After doing this, I like to use the function ```check_labels``` to make sure that the conversion was successful. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Convert training data into the DeepLabCut format\n",
    "import deeplabcut\n",
    "deeplabcut.convertcsv2h5(config_path, userfeedback=False)\n",
    "\n",
    "# Check labels (sanity check)\n",
    "deeplabcut.check_labels(config_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Labeled frame](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950852823-0AF8A8L896O58A94UXIV/094_093_089_val_AS_c2_frame325_bodypart.png?format=500w)\n",
    "\n",
    "*The labels appear to be in the right place.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structuring the project for testing\n",
    "\n",
    "Since bats are such a challenging animal to automatically digitize, in my lab, we've been relying on what I call \"refining\" to improve DLC accuracy. This just means that for each video that we want to analyze, we first digitize a few frames from it and include those frames as training data for the DeepLabCut network.\n",
    "\n",
    "We typically digitize approximately one frame per wingbeat. For a wingbeat frequency of 15 and a framerate of 800, this means approximately every 50th frame. For a lower framerate, say 400, we instead need to digitize every 50th frame. This is obviously less labour than digitizing every frame, but the workflow still scales poorly with increased acquisition. Therefore, I want to reduce the required amount of manual digitization. The challenge I've set myself is therefore this - **to use augmentation to match or beat the accuracy of refining**.\n",
    "\n",
    "### Test data\n",
    "I randomly chose seven digitized triplets of videos for testing. Every 50th frame of these will be used for testing.\n",
    "\n",
    "### Treatments\n",
    "I will use two \"treatments\" for testing. Firstly, I simply test how increasing the amount of training data affects accuracy. Secondly, I test how adding refining data affects accuracy. The result is a four permutations of treatments in total, sorted into four shuffles as shown in *table 1*.\n",
    "\n",
    "#### Refining\n",
    "The seven testing video triplets have had every 25th frame digitized. Frames 50, 100, 150, ... is used for testing. In this treatment, frames 25, 75, 175, ... will be used for used as training data, i.e. refining. I call the networks not trained on refining data OOD for out of domain and the ones traioned on refining data Ref for refining.\n",
    "\n",
    "#### Amount of base training data\n",
    "Subtracting the seven testing video triplets, I'm left with 19 (57 videos in total). I then randomly split these into two groups of 9 and 10. I will either train on only 10 of these video triplets, or on all 19. I call these treatments \"half\" and \"full\".\n",
    "\n",
    "This treatment might not seem relevant for figuring out how augmentation affects performance, and in truth, I chose to perform this test to answer another question, but you will see that the results of this test are very relevant to the challenge written in bold above.\n",
    "\n",
    "I will refer to training data from the 19 videos not tested on as \"base training data\".\n",
    "\n",
    "*Table 1: The two treatments result in four permutations, organized in four shuffles.*\n",
    "```\n",
    "|       |**Half**|**Full**|\n",
    "|-------|--------|--------|\n",
    "|**OOD**|shuffle1|shuffle2|\n",
    "|**Ref**|shuffle3|shuffle4|\n",
    "```\n",
    "\n",
    "### Training data indices\n",
    "By default, DeepLabCut assumes that we want to randomly divide the data into training and testing data. But for testing, you might want more control of this. Therefore, we will manually define which frames are for training and which are for testing. More on this later, but for now, we need to create a dummy training dataset as this will make DeepLabCut create a nice list of all the frames in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "\n",
    "# Dummy training dataset to get indexes and so on from later\n",
    "deeplabcut.create_training_dataset(config_path, Shuffles=[99])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define shuffles\n",
    "Now that we have created a dummy training dataset, we can parse its list of files and add the indices of the ones we want to train on to the list of training data indices, and the ones we want to test on to the test data indices.\n",
    "\n",
    "This is easy to do since when I converted our training data from DLTdv, I made sure to name the folders with endings indicating where they belong with regards to our treatments.\n",
    "\n",
    "The naming convention looks like this:\n",
    "- The folders that end with \"_A\" belong in the first half of the training data, that means that when we train on \"half\" the data, we only train on these.\n",
    "- The folders that end with \"_B\" belong in the second half of the training data, that means that when we train on the \"full\" data, we train on these as well as the ones from the first half.\n",
    "- The folders that end with \"_25_ref\" contain the refining frames from the 7 test videos, so when we train with the \"Ref\" data, we include these as training data.\n",
    "- The folders that end with \"_50_test\" are from the 7 test videos and they are always used as testing data.\n",
    "\n",
    "So by parsing the list of frames from the dataset, we can create our four shuffles as specified in *table 1* and in the list above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# create empty lists for putting testing and training indices in\n",
    "test_inds = []\n",
    "train_inds = []\n",
    "\n",
    "# train on half, test data is OOD, shuffle 1\n",
    "test_inds = []\n",
    "train_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_A\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "deeplabcut.create_training_dataset(\n",
    "    config_path,\n",
    "    Shuffles=[1],\n",
    "    trainIndices=[train_inds],\n",
    "    testIndices=[test_inds],\n",
    "    net_type=\"resnet_50\",\n",
    "    augmenter_type=\"../imagesaug\"\n",
    ")\n",
    "\n",
    "# train on half+ref, shuffle 2\n",
    "test_inds = []\n",
    "train_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_A\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_25_ref\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "deeplabcut.create_training_dataset(\n",
    "    config_path,\n",
    "    Shuffles=[2],\n",
    "    trainIndices=[train_inds],\n",
    "    testIndices=[test_inds],\n",
    "    net_type=\"resnet_50\",\n",
    "    augmenter_type=\"../imagesaug\"\n",
    ")\n",
    "\n",
    "# train on full, test data is OOD, shuffle 3\n",
    "test_inds = []\n",
    "train_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_A\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_B\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "deeplabcut.create_training_dataset(\n",
    "    config_path,\n",
    "    Shuffles=[3],\n",
    "    trainIndices=[train_inds],\n",
    "    testIndices=[test_inds],\n",
    "    net_type=\"resnet_50\",\n",
    "    augmenter_type=\"../imagesaug\"\n",
    ")\n",
    "\n",
    "# train on full+ref, shuffle 4\n",
    "test_inds = []\n",
    "train_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_A\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_B\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_25_ref\"):\n",
    "        train_inds.append(i)\n",
    "    elif str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "deeplabcut.create_training_dataset(\n",
    "    config_path,\n",
    "    Shuffles=[4],\n",
    "    trainIndices=[train_inds],\n",
    "    testIndices=[test_inds],\n",
    "    net_type=\"resnet_50\",\n",
    "    augmenter_type=\"../imagesaug\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using modelprefixes to keep track of augmentations\n",
    "\n",
    "Now we are using the shuffles to keep track of the training data used for our different networks, but we will also be trying different augmentatns. So how do we keep track of that? We could just keep making new shuffles and keep a list of what shuffle is what, but that can quickly become unwieldy.\n",
    "\n",
    "Instead, we can use modelprefixes. They provide a way to structure a DLC project into something like subprojects. Basically, if you copy a dlc-model folder (in my case, found in ```/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/dlc-models/iteration-0```) into a subfolder in the project, you can then interact with the copied dlc model by using the argument ```modelprefix```.\n",
    "\n",
    "Say, for example, that I want to try a more aggressive motion blur augmentation, then I could create a folder called ```extra_blurry``` in my project, copy the dlc-models into it, and then change the pose_cfg.yaml files that contain the augmentation parameters to increase the motion blur augmentation.\n",
    "\n",
    "I would then train the model with this command:\n",
    "```\n",
    "deeplabcut.train_network(\n",
    "    config_path,\n",
    "    shuffle=shuffle_number,\n",
    "    modelprefix='extra_blurry',\n",
    "    )\n",
    "```\n",
    "\n",
    "Our first networks will be trained with the default augmentation, but it's still useful to create a folder for them as we will want to copy the dlc-models in the root project folder for later augmentations, so it makes sense to keep those clean and do all training with model prefixes. \n",
    "\n",
    "### Regarding trainingsetindex\n",
    "Before we can create our subfolders, we need to spend a minute on the trainingset indices.\n",
    "\n",
    "We often use different splits between training and testing data. This is a number < 1 that sets the proportion of the training dataset that should be used for training and testing. The default is 0.95, meaning 95% is used for training, and 5% for testing. This can be either be specified in the config.yaml file or deduced automatically by DeepLabCut if we manually specify the training and testing indices (like we did). In the case of the latter, the split is indicated by the filenames of the the files in the training-dataset folder in the project folder. In my case this would be:\n",
    "```\n",
    "/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18\n",
    "```\n",
    "\n",
    "Let's look at the contents of that folder to deduce the training splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "files = os.listdir(\"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18\")\n",
    "print(*files,sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I've worked a bit on other stuff in this project, so there is some stuff here we need not consider (51 and 52)- But we see files related to shuffles 1 to 4. Let's focus on the .mat files:\n",
    "- bat_augmentation_austin_2020_bat_data_DLC80shuffle1.mat\n",
    "- bat_augmentation_austin_2020_bat_data_DLC84shuffle2.mat\n",
    "- bat_augmentation_austin_2020_bat_data_DLC88shuffle3.mat\n",
    "- bat_augmentation_austin_2020_bat_data_DLC89shuffle4.mat\n",
    "\n",
    "Not that \"DLC\" here refers to the scorer for the training data, so if you used the scorer name \"StephinMerritt\", your shuffle1 file name could be \"crocodiles_galloping_2023_StephinMerritt95shuffle1.mat\".\n",
    "\n",
    "Anyway, we see that there are numbers between the scorer name and the shuffle name, these represent the fraction of the training data compared to the training data plus the testing data.\n",
    "\n",
    "So for shuffle 1, of the data used, 80% is used as training data, for shuffle 2, that number is instead 84%, and so on. These numbers are called training fractions.\n",
    "\n",
    "In config.yaml, there is a property called ```TrainingFraction```. This is a list of the training fractions used in the project. In this project, that list needs to be changed to look like this:\n",
    "\n",
    "![TrainingFraction](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950855466-W5BOH6B4JJCO4M6RNVFG/config002_TrainingFraction.png?format=1000w)\n",
    "\n",
    "When we then reference a shuffle, we also need to reference the corresponding training fraction, this is used by specifying the parameter ```trainingsetindex```, and note that we start counting from zero. So for each of our shuffles, the index of the training fraction is one smaller than the shuffle number.\n",
    "\n",
    "With all that clarified, let's now create our first project subfolder/modelprefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = '/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "\n",
    "# we also need the package os for folder manipulation\n",
    "import os\n",
    "# and shutil for copying files\n",
    "import shutil\n",
    "\n",
    "#import tools for reading our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config\n",
    "\n",
    "# Number and name for our model folder\n",
    "model_number = 0\n",
    "modelprefix_pre = 'data_augm'\n",
    "daug_str = 'base'\n",
    "\n",
    "# Get config as dict and associated paths\n",
    "cfg = read_config(config_path)\n",
    "project_path = cfg[\"project_path\"] # or: os.path.dirname(config_path) #dlc_models_path = os.path.join(project_path, \"dlc-models\")\n",
    "training_datasets_path = os.path.join(project_path, \"training-datasets\")\n",
    "\n",
    "# Define shuffles\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0, 1, 2, 3]\n",
    "\n",
    "# Get train and test pose config file paths from base project, for each shuffle\n",
    "list_base_train_pose_config_file_paths = []\n",
    "list_base_test_pose_config_file_paths = []\n",
    "for shuffle_number, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    base_train_pose_config_file_path_TEMP,\\\n",
    "    base_test_pose_config_file_path_TEMP,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle_number,\n",
    "                                            trainingsetindex=trainingsetindex)  # base_train_pose_config_file\n",
    "    list_base_train_pose_config_file_paths.append(base_train_pose_config_file_path_TEMP)\n",
    "    list_base_test_pose_config_file_paths.append(base_test_pose_config_file_path_TEMP)\n",
    "\n",
    "# Create subdirs for this augmentation method\n",
    "model_prefix = '_'.join([modelprefix_pre, \"{0:0=2d}\".format(model_number), daug_str]) # modelprefix_pre = aug_\n",
    "aug_project_path = os.path.join(project_path, model_prefix)\n",
    "aug_dlc_models = os.path.join(aug_project_path, \"dlc-models\", )\n",
    "\n",
    "# make the folder for this modelprefix\n",
    "try:\n",
    "    os.mkdir(aug_project_path)\n",
    "except OSError as error:\n",
    "    print(error)\n",
    "    print(\"Skipping this one as it already exists\")\n",
    "\n",
    "# Copy base train pose config file to the directory of this augmentation method\n",
    "for j, (shuffle, trainingsetindex) in enumerate(zip(shuffles,trainingsetindices)):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    one_test_pose_config_file_path,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "    \n",
    "    # make train and test directories for this subdir\n",
    "    os.makedirs(str(os.path.dirname(one_train_pose_config_file_path))) # create parentdir 'train'\n",
    "    os.makedirs(str(os.path.dirname(one_test_pose_config_file_path))) # create parentdir 'test\n",
    "    \n",
    "    # copy test and train config from base project to this subdir\n",
    "    # copy base train config file\n",
    "    shutil.copyfile(list_base_train_pose_config_file_paths[j],\n",
    "                        one_train_pose_config_file_path) \n",
    "    # copy base test config file\n",
    "    shutil.copyfile(list_base_test_pose_config_file_paths[j],\n",
    "                        one_test_pose_config_file_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Edit pose_cfg.yaml\n",
    "We will also edit the settings for this network to use the adam optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "model_prefix = 'data_augm_00_base'\n",
    "\n",
    "## Initialise dict with additional edits to train config: optimizer\n",
    "train_edits_dict = {}\n",
    "dict_optimizer = {'optimizer':'adam',\n",
    "    'batch_size': 8, # the gpu I'm using has plenty of memory so batch size 8 makes sense\n",
    "    'multi_step': [[1e-4, 7500], [5 * 1e-5, 12000], [1e-5, 150000]]} # if no yaml file passed, initialise as an empty dict\n",
    "train_edits_dict.update({'optimizer': dict_optimizer['optimizer'], #'adam',\n",
    "    'batch_size': dict_optimizer['batch_size'],\n",
    "    'multi_step': dict_optimizer['multi_step']})\n",
    "\n",
    "for shuffle, trainingsetindex in zip(shuffles,trainingsetindices):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    _,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "\n",
    "    edit_config(str(one_train_pose_config_file_path), train_edits_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train networks\n",
    "Now we need to train our four shuffles. This will take several hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "# define config path and model prefix\n",
    "config_path='/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "model_prefix = 'data_augm_00_base'\n",
    "\n",
    "# the computer I'm working on has several gpus, here I used the third one.\n",
    "gputouse=3\n",
    "\n",
    "# define shuffles and trainingsetindices\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "# loop over shuffles and train each\n",
    "for shuffle, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    deeplabcut.train_network(\n",
    "        config_path,\n",
    "        shuffle=shuffle,\n",
    "        modelprefix=model_prefix,\n",
    "        gputouse=gputouse,\n",
    "        trainingsetindex=trainingsetindex,\n",
    "        max_snapshots_to_keep=3, # training for 150000 iterations so let's save 50, 100, and 150.\n",
    "        saveiters=50000\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Perform evaluation\n",
    "After training the networks, we need to test how they perform. To do that, we first need to evaluate them using the built-in DeepLabCut method ```evaluate_network```. By default this will evaluate the last snapshot, in our case, this means that the network will be evaluated after 150k iterations, but we want to know how well it does at 50k and 100k too, so let's edit our config file to test all saved snapshot.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = '/home/juser/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config, edit_config\n",
    "\n",
    "edit_config(config_path,{'snapshotindex':'all'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can evaluate the all three snapshots of all three shuffles. Remember that each shuffle has a different ```trainingsetindex```, so we need to perform the evaluations in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix = 'data_augm_00_base'\n",
    "Shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "for shuffle, trainingsetindex in zip(Shuffles,trainingsetindices):\n",
    "    deeplabcut.evaluate_network(config_path, modelprefix = model_prefix, Shuffles = [shuffle], trainingsetindex=trainingsetindex)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get data for plotting and statistical comparison\n",
    "The ```evaluate_network``` function calculates the euclidean distances between predicted points and groundtruth (i.e. human labels), but it only shows the average for that snapshot and shuffle over all the test data. What we want are long lists of the distances between ground truth points and predicted points. That way, on top of calculating the average, we can also perform stastistical testing.\n",
    "\n",
    "We will use a function I wrote called ```getErrorDistribution``` to get the list mentioned above. We will do this for each shuffle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix = 'data_augm_00_base'\n",
    "Shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') #my python file for this function is stored in the parent folder as I'm running this\n",
    "from getErrorDistribution import getErrorDistribution #import the getErrorDistribution function\n",
    "import numpy as np\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# get test indices\n",
    "test_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "error_distributions = []\n",
    "error_distributions_pcut = []\n",
    "\n",
    "for shuffle, trainFractionIndex in zip(Shuffles,trainingsetindices):\n",
    "   error_distributions_temp = []\n",
    "   error_distributions_pcut_temp = []\n",
    "   for snapshot in [0,1,2]: #we saved three snapshots, one at 50k iteratinos, one at 100k, and one at 150k\n",
    "      (\n",
    "            ErrorDistribution_all,\n",
    "            ErrorDistribution_test,\n",
    "            ErrorDistribution_train,\n",
    "            ErrorDistributionPCutOff_all,\n",
    "            _,\n",
    "            _\n",
    "      )  = getErrorDistribution(\n",
    "            config_path,\n",
    "            shuffle=shuffle,\n",
    "            snapindex=snapshot,\n",
    "            trainFractionIndex = trainFractionIndex,\n",
    "            modelprefix = model_prefix\n",
    "      )\n",
    "      error_distributions_temp.append(ErrorDistribution_all.iloc[test_inds].values.flatten())\n",
    "      error_distributions_pcut_temp.append(ErrorDistributionPCutOff_all.iloc[test_inds].values.flatten())\n",
    "   error_distributions.append(error_distributions_temp)\n",
    "   error_distributions_pcut.append(error_distributions_pcut_temp)\n",
    "error_distributionsb = np.array(error_distributions) # array with dimensions [shuffle, snapshot, frames]\n",
    "error_distributions_pcut = np.array(error_distributions_pcut) # array with dimensions [shuffle, snapshot, frames]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, (ax1,ax2) = plt.subplots(1,2)\n",
    "fig.set_figheight(8)\n",
    "fig.set_figwidth(10)\n",
    "\n",
    "for shuffle in [0,1,2,3]: #we start counting at 0, so for now, let's consider each index one less\n",
    "    ax1.errorbar(np.array([50,100,150])-1.5+shuffle,np.nanmean(error_distributions[shuffle,:],axis=1), np.nanstd(error_distributions[shuffle,:],axis=1)/len(test_inds)**.5)\n",
    "    ax2.errorbar(np.array([50,100,150])-1.5+shuffle,np.nanmean(error_distributions_pcut[shuffle,:],axis=1), np.nanstd(error_distributions_pcut[shuffle,:],axis=1)/len(test_inds)**.5)\n",
    "\n",
    "ax1.set_xticks([50, 100, 150])\n",
    "ax2.set_xticks([50, 100, 150])\n",
    "ax1.set_xlim([25,175])\n",
    "ax2.set_xlim([25,175])\n",
    "ax1.set_ylim([0,23])\n",
    "ax2.set_ylim([0,23])\n",
    "ax1.set_title(\"Without P-cut\")\n",
    "ax2.set_title(\"With P-cut 0.6\")\n",
    "ax2.set_yticklabels([])\n",
    "ax2.legend([\"half, OOD\",\"half, Ref\",\"full, OOD\", \"full, Ref\"])\n",
    "\n",
    "# add a big axis, hide frame\n",
    "fig.add_subplot(111, frameon=False)\n",
    "## hide tick and tick label of the big axis\n",
    "plt.tick_params(labelcolor='none', which='both', top=False, bottom=False, left=False, right=False)\n",
    "plt.xlabel(\"Iterations (thousands)\")\n",
    "plt.ylabel(\"Error (px)\")\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![download.png](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950856669-RVS6GMR4411BUVXY74FN/pre_aug_performance.png?format=750w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpret plot\n",
    "Just looking at the plot, a few things seem clear, let's start with the obvious and least surprising. By training a network on a subset of frames of the videos to be analyzed, we really improve accuracy. When doing this, i.e. when refining, increasing the amount of training data, i.e. frames not belonging to the videos to be analyzed, has little effect on accuracy.\n",
    "\n",
    "But surprisingly, when not refining, it appears that increasing the amount of training data can worsen performance. What's going on? We will get back to this later, but for now, let's just confirm my points here with some statistical testing.\n",
    "\n",
    "For the statistical comparison, we'll use a paired Wilcoxon ranked sum test, since we can't assume the data to be normally distributed. We'll do this for the set of errors with P-cut filter applied since doing it on both would be tedious and I assume most users will use some form of P-cut on their data. Furthermore, I will do the statistical comparison only on  the last snapshot since, again, doing it on all three would be tedious."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.stats import wilcoxon\n",
    "\n",
    "p_value = np.empty((4,4))\n",
    "p_value[:]=np.NaN\n",
    "\n",
    "for i in [0,1,2,3]:\n",
    "    for j in [0,1,2,3]:\n",
    "        if j<=i: continue\n",
    "        _, p_value[i,j] = wilcoxon(x = error_distributions_pcut[i,-1,:], y = error_distributions_pcut[j,-1,:],nan_policy='omit')\n",
    "\n",
    "print(p_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot per video\n",
    "The statistics are in line with expectations, except for regarding the difference between the two refined shuffles (shuffle2: \"half, ref\" and shuffle4: \"full, ref\"). The two of them being significantly different surprises me and makes me want to investigate the error distributions more in detail.\n",
    "\n",
    "To that end, let's plot the error distributions of the four shuffles per camera and per video number. Remember, out seven test video triplets each consist of three videos depicting the same bat flight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix = 'data_augm_00_base'\n",
    "Shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') #my python file for this function is stored in the parent folder as I'm running this\n",
    "from getErrorDistribution import getErrorDistribution #import the getErrorDistribution function\n",
    "import numpy as np\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# get test indices\n",
    "test_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "# this gives us the paths of our 27 test videos\n",
    "test_paths = list(set([image_paths[i][1] for i in test_inds]))\n",
    "\n",
    "#%% sorted so that the corresponding videos have the same index in three lists (one per camera)\n",
    "test_paths_cam1 =   ['TS5-544-Cam1_2020-06-25_000099Track8_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000103Track3_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000104Track3_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000108Track6_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000123Track6_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000128Track2_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000134Track5_50_test'\n",
    "                    ]\n",
    "test_paths_cam2 =   ['IL5-519-Cam2_2020-06-25_000099Track6_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000103Track3_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000104Track2_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000109Track1_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000124Track9_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000130Track2_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000136Track10_50_test'\n",
    "                    ]\n",
    "test_paths_cam3 =   ['IL5-534-Cam3_2020-06-25_000095Track14_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000100Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000101Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000106Track3_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000122Track7_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000127Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000133Track9_50_test'\n",
    "                    ]\n",
    "\n",
    "nvideos = 7 # number of videos\n",
    "\n",
    "# get test frame indexes per camera\n",
    "test_inds_cam1 = [[],[],[],[],[],[],[]]\n",
    "test_inds_cam2 = [[],[],[],[],[],[],[]]\n",
    "test_inds_cam3 = [[],[],[],[],[],[],[]]\n",
    "\n",
    "for i, path in enumerate(image_paths):\n",
    "    for j in range(nvideos):\n",
    "        if str(path[1]).__eq__(test_paths_cam1[j]):\n",
    "            test_inds_cam1[j].append(i)\n",
    "        elif str(path[1]).__eq__(test_paths_cam2[j]):\n",
    "            test_inds_cam2[j].append(i)\n",
    "        elif str(path[1]).__eq__(test_paths_cam3[j]):\n",
    "            test_inds_cam3[j].append(i)\n",
    "\n",
    "nshuffles = len(Shuffles)\n",
    "\n",
    "#pre-allocate matrixes for mean values and standard errors\n",
    "mean_cam1 = np.zeros([nshuffles,nvideos]) # shuffle x movie\n",
    "mean_cam2 = np.zeros([nshuffles,nvideos])\n",
    "mean_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "ste_cam1 = np.zeros([nshuffles,nvideos])\n",
    "ste_cam2 = np.zeros([nshuffles,nvideos])\n",
    "ste_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "meanPcut_cam1 = np.zeros([nshuffles,nvideos]) # shuffle x movie\n",
    "meanPcut_cam2 = np.zeros([nshuffles,nvideos])\n",
    "meanPcut_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "stePcut_cam1 = np.zeros([nshuffles,nvideos])\n",
    "stePcut_cam2 = np.zeros([nshuffles,nvideos])\n",
    "stePcut_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "# %%\n",
    "\n",
    "for i, shuffle in enumerate(Shuffles):    \n",
    "    trainFractionIndex = i\n",
    "    snapshot=-1\n",
    "    (\n",
    "        ErrorDistribution_all,\n",
    "        _,\n",
    "        _,\n",
    "        ErrorDistributionPCutOff_all,\n",
    "        _,\n",
    "        _\n",
    "    )  = getErrorDistribution(\n",
    "        config_path,\n",
    "        shuffle=shuffle,\n",
    "        snapindex=snapshot,\n",
    "        trainFractionIndex = trainFractionIndex,\n",
    "        modelprefix = model_prefix\n",
    "    )\n",
    "    for movie_number in range(7):\n",
    "\n",
    "        meanPcut_cam1[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])\n",
    "        stePcut_cam1[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam1[movie_number]][:].size**.5)\n",
    "\n",
    "        meanPcut_cam2[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam2[movie_number]][:])\n",
    "        stePcut_cam2[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam2[movie_number]][:].size**.5)\n",
    "\n",
    "        meanPcut_cam3[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam3[movie_number]][:])\n",
    "        stePcut_cam3[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam3[movie_number]][:].size**.5)\n",
    "\n",
    "fig, (ax1,ax2,ax3) = plt.subplots(3,1)\n",
    "fig.set_figheight(15)\n",
    "fig.set_figwidth(10)\n",
    "for i, shuffle in enumerate(Shuffles):\n",
    "    \n",
    "    # to jitter the error bars to keep  them from overlapping\n",
    "    movie_number = list(range(1,8))\n",
    "    movie_number = [x - 2/50 + shuffle/50 for x in movie_number]\n",
    "    \n",
    "    ax1.errorbar(movie_number,meanPcut_cam1[i,:], stePcut_cam1[i,:,])\n",
    "\n",
    "    ax2.errorbar(movie_number,meanPcut_cam2[i,:], stePcut_cam2[i,:,])\n",
    "\n",
    "    ax3.errorbar(movie_number,meanPcut_cam3[i,:], stePcut_cam3[i,:,])\n",
    "\n",
    "ax1.set_ylim([0, 50])\n",
    "ax2.set_ylim([0, 50])\n",
    "ax3.set_ylim([0, 50])\n",
    "ax1.set_title(\"Cam 1\")\n",
    "ax2.set_title(\"Cam 2\")\n",
    "ax3.set_title(\"Cam 3\")\n",
    "ax1.set_ylabel(\"Error (px\")\n",
    "ax2.set_ylabel(\"Error (px\")\n",
    "ax3.set_ylabel(\"Error (px\")\n",
    "ax1.legend([\"half, OOD\",\"half, Ref\",\"full, OOD\", \"full, Ref\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![download2.png](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950857302-7FYQ87AQ5OKGV6D8D9DF/pre_aug_per_video.png?format=500w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpret per video plot\n",
    "The first thing that becomes clear to me from inspecting this plot is that there are some videos on which all shuffles perform well no matter refining. These would be video 3, 6 and 7. The second thing is that there are instances where \"half, ref\" does better than \"full, ref\" and vice versa, which could explain why they are significantly different according to the statistical test performed above, despite their overall means being relatively similar according to the plot of the mean errors on the full error distributions.\n",
    "\n",
    "## Compare \"easy\" and \"difficult\" videos\n",
    "For shuffle3, \"full, OOD\", video 1, 2, 4, and 5 does poorly whereas video 3, 6, and 7 do rather well. Why is this? To get an idea of if the videos on which the network does well differs from the one on which it does poorly, we can create labeled frames with the predicted positions of the bodyparts and the correct positions (human labels). This is done with the ```evaluate_network``` function using the ```plotting``` input parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = '/home/jusers/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "model_prefix = 'data_augm_00_base'\n",
    "\n",
    "# we only want to plot the last snapshot (150k iterations)\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config, edit_config\n",
    "\n",
    "edit_config(config_path,{'snapshotindex':-1})\n",
    "\n",
    "shuffle = 3\n",
    "trainingsetindex = 2\n",
    "\n",
    "deeplabcut.evaluate_network(config_path, modelprefix = model_prefix, Shuffles = [shuffle], trainingsetindex=trainingsetindex, plotting=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect DeepLabCut predictions from evaluations\n",
    "\n",
    "Inspecting and comparing frames reveal an interesting relationship. It appears the network does well when the bat is flying left-right, and bad when the bat is flying right to left.\n",
    "\n",
    "The two following example images reveal what I mean.\n",
    "\n",
    "![\"bad\"](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950853195-EBZD23OVFO0NUSW94CUT/Test-IL5-519-Cam2_2020-06-25_000103Track3_50_test-103_103_100_AS_c2_frame200.png?format=500w)\n",
    "\n",
    "*The bat is flying from the right to the left, and the network has mistaken the left side of the bat for the right side.*\n",
    "\n",
    "![\"good\"](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950854190-2VRZO69FEZXJ7GZGFV6J/Test-IL5-519-Cam2_2020-06-25_000130Track2_50_test-128_130_127_AS_c2_frame300.png?format=500w)\n",
    "\n",
    "*The bat is flying from the left to the right, and the network is accurately predicting all points.\n",
    "\n",
    "The inspection of the predicted bodypart locations reveal that the low accuracy for some videos is caused by the network mixing up left and right. This could also explain why the performance goes down when we add more data. It appears that the added data results in the network becoming biased towards one flight direction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augment to reduce left-right bias\n",
    "Our evaluation showed that our network has a bias towards bats flying in one direction. In an attempt to reduce that bias and make the network better at telling left from right independent of flight direction, we will use a left-right flipping augmentaton.\n",
    "\n",
    "First up, let's make another modelprefix folder in our project. It basically the same code as before but with some minor modifications, I'll highlight the changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = '/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "\n",
    "# we also need the package os for folder manipulation\n",
    "import os\n",
    "# and shutil for copying files\n",
    "import shutil\n",
    "\n",
    "#import tools for reading our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config\n",
    "\n",
    "# Number and name for our model folder\n",
    "model_number = 1 # CHANGE\n",
    "modelprefix_pre = 'data_augm'\n",
    "daug_str = 'fliplr' # CHANGE\n",
    "\n",
    "# Get config as dict and associated paths\n",
    "cfg = read_config(config_path)\n",
    "project_path = cfg[\"project_path\"] # or: os.path.dirname(config_path) #dlc_models_path = os.path.join(project_path, \"dlc-models\")\n",
    "training_datasets_path = os.path.join(project_path, \"training-datasets\")\n",
    "\n",
    "# Define shuffles\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0, 1, 2, 3]\n",
    "\n",
    "# Get train and test pose config file paths from base project, for each shuffle\n",
    "list_base_train_pose_config_file_paths = []\n",
    "list_base_test_pose_config_file_paths = []\n",
    "for shuffle_number, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    base_train_pose_config_file_path_TEMP,\\\n",
    "    base_test_pose_config_file_path_TEMP,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle_number,\n",
    "                                            trainingsetindex=trainingsetindex)  # base_train_pose_config_file\n",
    "    list_base_train_pose_config_file_paths.append(base_train_pose_config_file_path_TEMP)\n",
    "    list_base_test_pose_config_file_paths.append(base_test_pose_config_file_path_TEMP)\n",
    "\n",
    "# Create subdirs for this augmentation method\n",
    "model_prefix = '_'.join([modelprefix_pre, \"{0:0=2d}\".format(model_number), daug_str]) # modelprefix_pre = aug_\n",
    "aug_project_path = os.path.join(project_path, model_prefix)\n",
    "aug_dlc_models = os.path.join(aug_project_path, \"dlc-models\", )\n",
    "\n",
    "# make the folder for this modelprefix\n",
    "try:\n",
    "    os.mkdir(aug_project_path)\n",
    "except OSError as error:\n",
    "    print(error)\n",
    "    print(\"Skipping this one as it already exists\")\n",
    "\n",
    "# Copy base train pose config file to the directory of this augmentation method\n",
    "for j, (shuffle, trainingsetindex) in enumerate(zip(shuffles,trainingsetindices)):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    one_test_pose_config_file_path,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "    \n",
    "    # make train and test directories for this subdir\n",
    "    os.makedirs(str(os.path.dirname(one_train_pose_config_file_path))) # create parentdir 'train'\n",
    "    os.makedirs(str(os.path.dirname(one_test_pose_config_file_path))) # create parentdir 'test\n",
    "    \n",
    "    # copy test and train config from base project to this subdir\n",
    "    # copy base train config file\n",
    "    shutil.copyfile(list_base_train_pose_config_file_paths[j],\n",
    "                        one_train_pose_config_file_path) \n",
    "    # copy base test config file\n",
    "    shutil.copyfile(list_base_test_pose_config_file_paths[j],\n",
    "                        one_test_pose_config_file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Edit pose_cfg.yaml for fliplr augmentation\n",
    "As before, we will edit the settings for this network to use the adam optimizer. More importantly, though, we will apply the ```fliplr``` augmentation.\n",
    "\n",
    "This randomly flips some training frames left and right during training. Now, this alone would lead to worse performance, since we would be training the network to think that the left wing is the right wing, and vice versa. But this augmentation method takes that into account, and in addition to mirroring the image, it also swaps the labels from left to right and vice versa. To accomlish this, we need to supply the pose_cfg.yaml file with a list of which points should be swapped.\n",
    "\n",
    "The image of a bat outline above shows he locations of the labels/bodyparts on the bats, and from it, it's clear which points should be swapped, all points that aren't on the bilateral have a name and end in eithr R or L. So t3R should be swapped with t3L when flipping, and so on.\n",
    "\n",
    "The labels are these, and in this order:\n",
    "0. 't3L'\n",
    "1. 'wstL'\n",
    "2. 't5L'\n",
    "3. 'elbL'\n",
    "4. 'shdL'\n",
    "5. 'ankL'\n",
    "6. 'nl'\n",
    "7. 'str'\n",
    "8. 'lmb'\n",
    "9. 'shdR'\n",
    "10. 'ankR'\n",
    "11. 'elbR'\n",
    "12. 'wstR'\n",
    "13. 't5R'\n",
    "14. 't3R'\n",
    "15. 'tail'\n",
    "\n",
    "So the symmetric pairs are the following:\n",
    "- t3:   (0, 14)\n",
    "- wst:  (1, 12)\n",
    "- t5:   (2, 13)\n",
    "- elb:  (3, 11)\n",
    "- shd:  (4, 9)\n",
    "- ank:  (5, 10)\n",
    "\n",
    "Now, let's edit the ```pose_cfg.yaml``` accordingly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "#import tools for changing our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import edit_config\n",
    "\n",
    "model_prefix = 'data_augm_01_fliplr'\n",
    "\n",
    "## Initialise dict with additional edits to train config: optimizer\n",
    "train_edits_dict = {}\n",
    "dict_optimizer = {'optimizer':'adam',\n",
    "    'batch_size': 8, # the gpu I'm using has plenty of memory so batch size 8 makes sense\n",
    "    'multi_step': [[1e-4, 7500], [5 * 1e-5, 12000], [1e-5, 150000]]} # if no yaml file passed, initialise as an empty dict\n",
    "train_edits_dict.update({'optimizer': dict_optimizer['optimizer'], #'adam',\n",
    "    'batch_size': dict_optimizer['batch_size'],\n",
    "    'multi_step': dict_optimizer['multi_step']})\n",
    "\n",
    "# Augmentation edits\n",
    "edits_dict = dict()\n",
    "edits_dict[\"symmetric_pairs\"] = (0, 14), (1, 12), (2, 13), (3, 11), (4, 9), (5, 10)\n",
    "edits_dict[\"fliplr\"] = True\n",
    "\n",
    "for shuffle, trainingsetindex in zip(shuffles,trainingsetindices):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    _,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "\n",
    "    edit_config(str(one_train_pose_config_file_path), edits_dict)\n",
    "    edit_config(str(one_train_pose_config_file_path), train_edits_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train network with ```fliplr``` augmentation\n",
    "Now let's train the network with the new augmentation. Again, this will take several hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "# define config path and model prefix\n",
    "config_path='/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "model_prefix = 'data_augm_01_flipr'\n",
    "\n",
    "# the computer I'm working on has several gpus, here I used the third one.\n",
    "gputouse=3\n",
    "\n",
    "# define shuffles and trainingsetindices\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "# loop over shuffles and train each\n",
    "for shuffle, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    deeplabcut.train_network(\n",
    "        config_path,\n",
    "        shuffle=shuffle,\n",
    "        modelprefix=model_prefix,\n",
    "        gputouse=gputouse,\n",
    "        trainingsetindex=trainingsetindex,\n",
    "        max_snapshots_to_keep=3, # training for 150000 iterations so let's save 50, 100, and 150.\n",
    "        saveiters=50000\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare performance of fliplr augmented network to refined baseline\n",
    "\n",
    "Since the challenge I set is **to use augmentation to match or beat the accuracy of refining**, let's stick to comparing baseline \"Full, ref\" to fliplr augmented \"Full, OOD\".\n",
    "\n",
    "The process is similar to what we've already done, but now we are only comparing two networks.\n",
    "\n",
    "We will use ```evaluate_network``` to get create accuracy data, then we will use ```getErrorDistribution``` to access that data and plot the average performance with error bars. Same as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix = 'data_augm_01_fliplr'\n",
    "Shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "#import tools for modifying our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import edit_config\n",
    "\n",
    "# make sure we are testing all snapshots\n",
    "edit_config(config_path,{'snapshotindex':'all'})\n",
    "\n",
    "for shuffle, trainingsetindex in zip(Shuffles,trainingsetindices):\n",
    "    deeplabcut.evaluate_network(config_path, modelprefix = model_prefix, Shuffles = [shuffle], trainingsetindex=trainingsetindex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix_base = 'data_augm_00_base'\n",
    "model_prefix_augm = 'data_augm_01_fliplr'\n",
    "Shuffles = [4,3] # let's start with the refined un-augmented, i.e. shuffle 4\n",
    "trainingsetindices = [3,2]\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') #my python file for this function is stored in the parent folder as I'm running this\n",
    "from getErrorDistribution import getErrorDistribution #import the getErrorDistribution function\n",
    "import numpy as np\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# get test indices\n",
    "test_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "error_distributions_pcut = []\n",
    "\n",
    "for shuffle, trainFractionIndex in zip(Shuffles,trainingsetindices):\n",
    "   error_distributions_pcut_temp = []\n",
    "   if shuffle == 4: model_prefix = model_prefix_base\n",
    "   elif shuffle == 3: model_prefix = model_prefix_augm\n",
    "   \n",
    "   for snapshot in [0,1,2]: #we saved three snapshots, one at 50k iteratinos, one at 100k, and one at 150k\n",
    "      (\n",
    "            _,\n",
    "            _,\n",
    "            _,\n",
    "            ErrorDistributionPCutOff_all,\n",
    "            _,\n",
    "            _\n",
    "      )  = getErrorDistribution(\n",
    "            config_path,\n",
    "            shuffle=shuffle,\n",
    "            snapindex=snapshot,\n",
    "            trainFractionIndex = trainFractionIndex,\n",
    "            modelprefix = model_prefix\n",
    "      )\n",
    "      error_distributions_pcut_temp.append(ErrorDistributionPCutOff_all.iloc[test_inds].values.flatten())\n",
    "   error_distributions_pcut.append(error_distributions_pcut_temp)\n",
    "\n",
    "error_distributions_pcut = np.array(error_distributions_pcut) # array with dimensions [shuffle, snapshot, frames]\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(10, 5))\n",
    "for shuffle in [0,1]: #we start counting at 0, so for now, let's consider each index one less\n",
    "    plt.errorbar(np.array([50,100,150])-1.5+shuffle,np.nanmean(error_distributions_pcut[shuffle,:],axis=1), np.nanstd(error_distributions_pcut[shuffle,:],axis=1)/len(test_inds)**.5)\n",
    "\n",
    "plt.xticks([50, 100, 150])\n",
    "plt.xlim([25,175])\n",
    "plt.ylim([0,10])\n",
    "plt.title(\"Error with P-cut 0.6, comparing baseline to fliplr and 180 degrees rotation augmented\")\n",
    "\n",
    "plt.legend([\"Full, ref, baseline\", \"Full, OOD, fliplr\"])\n",
    "\n",
    "plt.xlabel(\"Iterations (thousands)\")\n",
    "plt.ylabel(\"Error (px)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpret performance of network after \"fliplr\" augmentation\n",
    "\n",
    "That is a significant improvement in accuracy. The ```fliplr``` augmentation seems to prevent the network from becoming biased to one flight direction. It is however, not quite beating the performance of the baseline refined case (shuffle4: full, Ref).\n",
    "\n",
    "There is another augmentation that, in our case, is related to the flight direction, namely the rotation augmentation. The idea here is that in order to prevent the network from associating absolute in-image spatial orientation with positions for certain labels, we want to randomly rotate some of the training data. The degree to which this makes sense is context dependent. If you are analyzing a horse walking in a horizontal direction from the camera's point of view, large degrees of rotation make little sense as the network will never encounter a horse walking upside down. But if you are filming an animal from above, or below if the animal is flying or swimming, such that the animal is free to move in any direction, then this type of augmentation makes sense.\n",
    "\n",
    "By default, DeepLabCut applies a 25° rotation augmentation, I tried increasing that, first to 90°, then to 180°. I will only cover the 180° case here as the 90° case is exactly the same, only with slightly worse performance.\n",
    "\n",
    "All the steps taken are the same as before.\n",
    "- Create the model prefix folder\n",
    "- Change the ```pose_cfg.yaml``` file to apply the augmentation\n",
    "- Train the network\n",
    "- Evaluate the network performance\n",
    "- Plot results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = '/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "\n",
    "# we also need the package os for folder manipulation\n",
    "import os\n",
    "# and shutil for copying files\n",
    "import shutil\n",
    "\n",
    "#import tools for reading our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import read_config\n",
    "\n",
    "# Number and name for our model folder\n",
    "model_number = 3 # CHANGE\n",
    "modelprefix_pre = 'data_augm'\n",
    "daug_str = 'max_rotate' # CHANGE\n",
    "\n",
    "# Get config as dict and associated paths\n",
    "cfg = read_config(config_path)\n",
    "project_path = cfg[\"project_path\"] # or: os.path.dirname(config_path) #dlc_models_path = os.path.join(project_path, \"dlc-models\")\n",
    "training_datasets_path = os.path.join(project_path, \"training-datasets\")\n",
    "\n",
    "# Define shuffles\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0, 1, 2, 3]\n",
    "\n",
    "# Get train and test pose config file paths from base project, for each shuffle\n",
    "list_base_train_pose_config_file_paths = []\n",
    "list_base_test_pose_config_file_paths = []\n",
    "for shuffle_number, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    base_train_pose_config_file_path_TEMP,\\\n",
    "    base_test_pose_config_file_path_TEMP,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle_number,\n",
    "                                            trainingsetindex=trainingsetindex)  # base_train_pose_config_file\n",
    "    list_base_train_pose_config_file_paths.append(base_train_pose_config_file_path_TEMP)\n",
    "    list_base_test_pose_config_file_paths.append(base_test_pose_config_file_path_TEMP)\n",
    "\n",
    "# Create subdirs for this augmentation method\n",
    "model_prefix = '_'.join([modelprefix_pre, \"{0:0=2d}\".format(model_number), daug_str]) # modelprefix_pre = aug_\n",
    "aug_project_path = os.path.join(project_path, model_prefix)\n",
    "aug_dlc_models = os.path.join(aug_project_path, \"dlc-models\", )\n",
    "\n",
    "# make the folder for this modelprefix\n",
    "try:\n",
    "    os.mkdir(aug_project_path)\n",
    "except OSError as error:\n",
    "    print(error)\n",
    "    print(\"Skipping this one as it already exists\")\n",
    "\n",
    "# Copy base train pose config file to the directory of this augmentation method\n",
    "for j, (shuffle, trainingsetindex) in enumerate(zip(shuffles,trainingsetindices)):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    one_test_pose_config_file_path,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "    \n",
    "    # make train and test directories for this subdir\n",
    "    os.makedirs(str(os.path.dirname(one_train_pose_config_file_path))) # create parentdir 'train'\n",
    "    os.makedirs(str(os.path.dirname(one_test_pose_config_file_path))) # create parentdir 'test\n",
    "    \n",
    "    # copy test and train config from base project to this subdir\n",
    "    # copy base train config file\n",
    "    shutil.copyfile(list_base_train_pose_config_file_paths[j],\n",
    "                        one_train_pose_config_file_path) \n",
    "    # copy base test config file\n",
    "    shutil.copyfile(list_base_test_pose_config_file_paths[j],\n",
    "                        one_test_pose_config_file_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "#import tools for changing our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import edit_config\n",
    "\n",
    "model_prefix = 'data_augm_03_max_rotate'\n",
    "\n",
    "## Initialise dict with additional edits to train config: optimizer\n",
    "train_edits_dict = {}\n",
    "dict_optimizer = {'optimizer':'adam',\n",
    "    'batch_size': 8, # the gpu I'm using has plenty of memory so batch size 8 makes sense\n",
    "    'multi_step': [[1e-4, 7500], [5 * 1e-5, 12000], [1e-5, 150000]]} # if no yaml file passed, initialise as an empty dict\n",
    "train_edits_dict.update({'optimizer': dict_optimizer['optimizer'], #'adam',\n",
    "    'batch_size': dict_optimizer['batch_size'],\n",
    "    'multi_step': dict_optimizer['multi_step']})\n",
    "\n",
    "# Augmentation edits\n",
    "edits_dict = dict()\n",
    "edits_dict[\"symmetric_pairs\"] = (0, 14), (1, 12), (2, 13), (3, 11), (4, 9), (5, 10)\n",
    "edits_dict[\"fliplr\"] = True\n",
    "edits_dict[\"rotation\"] = 180\n",
    "\n",
    "for shuffle, trainingsetindex in zip(shuffles,trainingsetindices):\n",
    "    one_train_pose_config_file_path,\\\n",
    "    _,\\\n",
    "    _ = deeplabcut.return_train_network_path(config_path,\n",
    "                                            shuffle=shuffle,\n",
    "                                            trainingsetindex=trainingsetindex,\n",
    "                                            modelprefix=model_prefix)\n",
    "\n",
    "    edit_config(str(one_train_pose_config_file_path), edits_dict)\n",
    "    edit_config(str(one_train_pose_config_file_path), train_edits_dict)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train network with ```fliplr``` and stronger ```rotation``` augmentation\n",
    "Now let's train the network with the new augmentation. Again, this will take several hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import deeplabcut\n",
    "# define config path and model prefix\n",
    "config_path='/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml'\n",
    "model_prefix = 'data_augm_03_max_rotate'\n",
    "\n",
    "# the computer I'm working on has several gpus, here I used the third one.\n",
    "gputouse=3\n",
    "\n",
    "# define shuffles and trainingsetindices\n",
    "shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "# loop over shuffles and train each\n",
    "for shuffle, trainingsetindex in zip(shuffles, trainingsetindices):\n",
    "    deeplabcut.train_network(\n",
    "        config_path,\n",
    "        shuffle=shuffle,\n",
    "        modelprefix=model_prefix,\n",
    "        gputouse=gputouse,\n",
    "        trainingsetindex=trainingsetindex,\n",
    "        max_snapshots_to_keep=3, # training for 150000 iterations so let's save 50, 100, and 150.\n",
    "        saveiters=50000\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Evaluate the network\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix = 'data_augm_03_max_rotate'\n",
    "Shuffles = [1,2,3,4]\n",
    "trainingsetindices = [0,1,2,3]\n",
    "\n",
    "#import tools for modifying our config file\n",
    "from deeplabcut.utils.auxiliaryfunctions import edit_config\n",
    "\n",
    "# make sure we are testing all snapshots\n",
    "edit_config(config_path,{'snapshotindex':'all'})\n",
    "\n",
    "for shuffle, trainingsetindex in zip(Shuffles,trainingsetindices):\n",
    "    deeplabcut.evaluate_network(config_path, modelprefix = model_prefix, Shuffles = [shuffle], trainingsetindex=trainingsetindex, gputouse=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Get detailed error data for plotting\n",
    "\n",
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix_base = 'data_augm_00_base'\n",
    "model_prefix_augm = 'data_augm_03_max_rotate'\n",
    "Shuffles = [4,3] # let's start with the refined un-augmented, i.e. shuffle 4\n",
    "trainingsetindices = [3,2]\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') #my python file for this function is stored in the parent folder as I'm running this\n",
    "from getErrorDistribution import getErrorDistribution #import the getErrorDistribution function\n",
    "import numpy as np\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# get test indices\n",
    "test_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "error_distributions_pcut = []\n",
    "\n",
    "for shuffle, trainFractionIndex in zip(Shuffles,trainingsetindices):\n",
    "   error_distributions_pcut_temp = []\n",
    "   if shuffle == 4: model_prefix = model_prefix_base\n",
    "   elif shuffle == 3: model_prefix = model_prefix_augm\n",
    "   \n",
    "   for snapshot in [0,1,2]: #we saved three snapshots, one at 50k iteratinos, one at 100k, and one at 150k\n",
    "      (\n",
    "            _,\n",
    "            _,\n",
    "            _,\n",
    "            ErrorDistributionPCutOff_all,\n",
    "            _,\n",
    "            _\n",
    "      )  = getErrorDistribution(\n",
    "            config_path,\n",
    "            shuffle=shuffle,\n",
    "            snapindex=snapshot,\n",
    "            trainFractionIndex = trainFractionIndex,\n",
    "            modelprefix = model_prefix\n",
    "      )\n",
    "      error_distributions_pcut_temp.append(ErrorDistributionPCutOff_all.iloc[test_inds].values.flatten())\n",
    "   error_distributions_pcut.append(error_distributions_pcut_temp)\n",
    "\n",
    "error_distributions_pcut = np.array(error_distributions_pcut) # array with dimensions [shuffle, snapshot, frames]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "for shuffle in [0,1]: #we start counting at 0, so for now, let's consider each index one less\n",
    "    plt.errorbar(np.array([50,100,150])-1.5+shuffle,np.nanmean(error_distributions_pcut[shuffle,:],axis=1), np.nanstd(error_distributions_pcut[shuffle,:],axis=1)/len(test_inds)**.5)\n",
    "\n",
    "plt.xticks([50, 100, 150])\n",
    "plt.xlim([25,175])\n",
    "plt.ylim([0,10])\n",
    "plt.title(\"Error with P-cut 0.6, comparing baseline to fliplr and 180 degrees rotation augmented\")\n",
    "\n",
    "plt.legend([\"Full, ref, baseline\", \"Full, OOD, fliplr_180_rotate\"])\n",
    "\n",
    "plt.xlabel(\"Iterations (thousands)\")\n",
    "plt.ylabel(\"Error (px)\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![download3.png](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1661950856089-7I0IXOS1G0D2YUJ9GNSB/flip_rot180.png?format=1500w)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this, it appears that we have reached our goal, which was **to use augmentation to match or beat the accuracy of refining**. But let's also take a look at the accuracy per video. I will also plot the baseline, un-refined full network accuracy here to highlight how much the augmentation has improved the performance of the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# in case we restarted the kernel or something, let's make\n",
    "# sure deeplabcut is imported and the config_path defined\n",
    "import deeplabcut\n",
    "\n",
    "config_path = \"/home/user/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/config.yaml\"\n",
    "model_prefix_base = 'data_augm_00_base'\n",
    "model_prefix_augm = 'data_augm_03_max_rotate'\n",
    "Shuffles = [4,3,3]\n",
    "trainingsetindices = [3,2,2]\n",
    "\n",
    "# We need pandas for creatig a nice list to parse\n",
    "import pandas as pd\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') #my python file for this function is stored in the parent folder as I'm running this\n",
    "from getErrorDistribution import getErrorDistribution #import the getErrorDistribution function\n",
    "import numpy as np\n",
    "\n",
    "# Read the h5 file containing all the frames, (project_folder/training-datasets/iteration-0/UnaufmentedDataSet_project_folder/CollectedData_LabelerName.h5)\n",
    "df = pd.read_hdf('/home/juser/projects/bat_augmentation_austin_2020_bat_data-DLC-2022-08-18/training-datasets/iteration-0/UnaugmentedDataSet_bat_augmentation_austin_2020_bat_dataAug18/CollectedData_DLC.h5')\n",
    "\n",
    "image_paths = df.index.to_list() # turn dataframe into list\n",
    "\n",
    "# get test indices\n",
    "test_inds = []\n",
    "for i, path in enumerate(image_paths):\n",
    "    if str(path[1]).endswith(\"_50_test\"):\n",
    "        test_inds.append(i)\n",
    "\n",
    "# this gives us the paths of our 27 test videos\n",
    "test_paths = list(set([image_paths[i][1] for i in test_inds]))\n",
    "\n",
    "#%% sorted so that the corresponding videos have the same index in three lists (one per camera)\n",
    "test_paths_cam1 =   ['TS5-544-Cam1_2020-06-25_000099Track8_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000103Track3_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000104Track3_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000108Track6_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000123Track6_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000128Track2_50_test',\n",
    "                    'TS5-544-Cam1_2020-06-25_000134Track5_50_test'\n",
    "                    ]\n",
    "test_paths_cam2 =   ['IL5-519-Cam2_2020-06-25_000099Track6_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000103Track3_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000104Track2_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000109Track1_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000124Track9_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000130Track2_50_test',\n",
    "                    'IL5-519-Cam2_2020-06-25_000136Track10_50_test'\n",
    "                    ]\n",
    "test_paths_cam3 =   ['IL5-534-Cam3_2020-06-25_000095Track14_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000100Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000101Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000106Track3_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000122Track7_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000127Track4_50_test',\n",
    "                    'IL5-534-Cam3_2020-06-25_000133Track9_50_test'\n",
    "                    ]\n",
    "\n",
    "nvideos = 7 # number of videos\n",
    "\n",
    "# get test frame indexes per camera\n",
    "test_inds_cam1 = [[],[],[],[],[],[],[]]\n",
    "test_inds_cam2 = [[],[],[],[],[],[],[]]\n",
    "test_inds_cam3 = [[],[],[],[],[],[],[]]\n",
    "\n",
    "for i, path in enumerate(image_paths):\n",
    "    for j in range(nvideos):\n",
    "        if str(path[1]).__eq__(test_paths_cam1[j]):\n",
    "            test_inds_cam1[j].append(i)\n",
    "        elif str(path[1]).__eq__(test_paths_cam2[j]):\n",
    "            test_inds_cam2[j].append(i)\n",
    "        elif str(path[1]).__eq__(test_paths_cam3[j]):\n",
    "            test_inds_cam3[j].append(i)\n",
    "\n",
    "nshuffles = len(Shuffles)\n",
    "\n",
    "#pre-allocate matrixes for mean values and standard errors\n",
    "mean_cam1 = np.zeros([nshuffles,nvideos]) # shuffle x movie\n",
    "mean_cam2 = np.zeros([nshuffles,nvideos])\n",
    "mean_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "ste_cam1 = np.zeros([nshuffles,nvideos])\n",
    "ste_cam2 = np.zeros([nshuffles,nvideos])\n",
    "ste_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "meanPcut_cam1 = np.zeros([nshuffles,nvideos]) # shuffle x movie\n",
    "meanPcut_cam2 = np.zeros([nshuffles,nvideos])\n",
    "meanPcut_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "stePcut_cam1 = np.zeros([nshuffles,nvideos])\n",
    "stePcut_cam2 = np.zeros([nshuffles,nvideos])\n",
    "stePcut_cam3 = np.zeros([nshuffles,nvideos])\n",
    "\n",
    "# %%\n",
    "\n",
    "for i, shuffle in enumerate(Shuffles):\n",
    "    if shuffle == 4 or i == 2: model_prefix = model_prefix_base\n",
    "    elif shuffle == 3: model_prefix = model_prefix_augm \n",
    "\n",
    "    trainFractionIndex = shuffle-1\n",
    "    snapshot=-1\n",
    "    (\n",
    "        ErrorDistribution_all,\n",
    "        _,\n",
    "        _,\n",
    "        ErrorDistributionPCutOff_all,\n",
    "        _,\n",
    "        _\n",
    "    )  = getErrorDistribution(\n",
    "        config_path,\n",
    "        shuffle=shuffle,\n",
    "        snapindex=snapshot,\n",
    "        trainFractionIndex = trainFractionIndex,\n",
    "        modelprefix = model_prefix\n",
    "    )\n",
    "    for movie_number in range(7):\n",
    "\n",
    "        meanPcut_cam1[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])\n",
    "        stePcut_cam1[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam1[movie_number]][:].size**.5)\n",
    "\n",
    "        meanPcut_cam2[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam2[movie_number]][:])\n",
    "        stePcut_cam2[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam2[movie_number]][:].size**.5)\n",
    "\n",
    "        meanPcut_cam3[i,movie_number] = np.nanmean(ErrorDistributionPCutOff_all.values[test_inds_cam3[movie_number]][:])\n",
    "        stePcut_cam3[i,movie_number] = np.nanstd(ErrorDistributionPCutOff_all.values[test_inds_cam1[movie_number]][:])/(ErrorDistribution_all.values[test_inds_cam3[movie_number]][:].size**.5)\n",
    "\n",
    "fig, (ax1,ax2,ax3) = plt.subplots(3,1)\n",
    "fig.set_figheight(15)\n",
    "fig.set_figwidth(10)\n",
    "for i, shuffle in enumerate(Shuffles):\n",
    "    \n",
    "    # to jitter the error bars to keep  them from overlapping\n",
    "    movie_number = list(range(1,8))\n",
    "    movie_number = [x - 2/50 + shuffle/50 for x in movie_number]\n",
    "    \n",
    "    ax1.errorbar(movie_number,meanPcut_cam1[i,:], stePcut_cam1[i,:,])\n",
    "\n",
    "    ax2.errorbar(movie_number,meanPcut_cam2[i,:], stePcut_cam2[i,:,])\n",
    "\n",
    "    ax3.errorbar(movie_number,meanPcut_cam3[i,:], stePcut_cam3[i,:,])\n",
    "\n",
    "ax1.set_ylim([0, 50])\n",
    "ax2.set_ylim([0, 50])\n",
    "ax3.set_ylim([0, 50])\n",
    "ax1.set_title(\"Cam 1\")\n",
    "ax2.set_title(\"Cam 2\")\n",
    "ax3.set_title(\"Cam 3\")\n",
    "ax1.set_ylabel(\"Error (px\")\n",
    "ax2.set_ylabel(\"Error (px\")\n",
    "ax3.set_ylabel(\"Error (px\")\n",
    "ax1.legend([\"Full, ref, baseline\", \"Full, OOD, fliplr_180_rotate\", \"Full, OOD, baseline\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Interpret performance of network after ```fliplr``` and 180° rotation augmentation\n",
    "The results look very promising. Video 1 keeps being difficult. This is likely due to this video depicting a bat flying in a manner that is atypical for the dataset used here. Most videos in the dataset depict rather straight-forward flybys, but video 1 shows a bat slowly ascending while turning. Even so, the augmentation has lowered the pixel error for video 1 compared to the baseline, and for all other videos, the augmentation has led to performance close to, or better than, the refined case.\n",
    "\n",
    "# Conclusion\n",
    "What follows are some of my conclusions from using augmentation to improve performance when tracking wing movement of bats using DeepLabCut.\n",
    "\n",
    "## Where scalability is not important, consider refining\n",
    "We saw from our baseline accuracy tests that when refining, there was no gain from increasing the amount of base training data. In practical terms this means that if you know which videos you want to analyze, and when you do not plan on using the network for analyzing other videos in the future, then it makes sense to prioritize digitizing frames from the videos to be analyzed.\n",
    "\n",
    "## Inspect bodypart predictions and augment with purpose\n",
    "At first, I tried a couple of permutations of different image augmentations and saw no or little performance improvement. I tried augmentations related to scale, brightness, blurriness, and probably some other that slip my mind at the moment. I more or less concluded that augmentations have little effect and that for a given set of training and test data, DeepLabCut's default training parameters are close to optimal and little can be gained by augmenting the training data. But then I looked closer at the labeled frames from ```evaluate_network``` and realized that the high pixel errors were mainly caused by the network getting left and right mixed up and augmented accordingly. Doing that, I saw an average accuracy improvement of more than 10 pixels. For this dataset, 10 pixels is quite big, for comparison, on the largest image in the test dataset, the forearm of the bat is less than 70 pixels long.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python [conda env:DEEPLABCUT_newGUI] *",
   "language": "python",
   "name": "conda-env-DEEPLABCUT_newGUI-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "d59ffae89e6bb4db73216b2b6c5dd5d3dc066a6f8df0277b318288143e8a2adf"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
